import datetime

import matplotlib.pyplot as plt
import numpy as np
from docx.shared import Cm
from docxtpl import DocxTemplate, InlineImage
from scipy.integrate import odeint

# Задача 1

# Определение параметров
EI = 9000  # Изгибная жёсткость [N*m^2]
P = 3000  # Вес груза [N]
l1 = 1.8  # Длина пролёта между опорами [m]
l2 = 0.9  # Длина консоли [m]
g = 9.81  # Ускорение свободного падения [m/s^2]

# Определение опорных реакций
Ra = P * l2 / l1
Rb = P * (l1 + l2) / l1


# Определение функции Хевисайда h(x-l1)
def heaviside(x: np.ndarray) -> list:
    """Единичная функция Хевисайда, которая ставит в соответствие вектору, полученному
      на вход, вектор того же размера, который состоит из нулей и единиц.

    Parameters
    ----------
    x : np.ndarray
        Вектор исходных значений

    Returns
    -------
    list
        Вектор нулей и единиц
    """
    return [1 if i > l1 else 0 for i in x]


x = np.linspace(0, l1 + l2, 50)
y = np.array(
    (
        P * (l1 + l2) / (6 * l1) * (x - l1) ** 3 * heaviside(x)
        - P * l2 / (6 * l1) * x**3
        + P * l2 / (6 * l1) * x
    )
    / EI
)

# Отображаем форму изогнутой оси
fig3 = plt.figure()
plt.plot(x, y, label="y(x)")
plt.xlabel("x")
plt.ylabel("y")
plt.title("Решение дифференциального уравнения с граничными условиями")
plt.legend()
plt.show()
fig3.savefig("dz/oscillation_analysis/images/image_bending.png")

# Вычисляем собственную частоту колебаний груза, используя значение прогиба балки при
#  x = l1 + l2

Natural_frequency = 1 / (2 * np.pi) * np.sqrt(g / abs(y[-1]))

# ЗАДАЧА 2


# Параметры системы
L1 = 1  # Длина маятника [m]
L2 = 1  # Длина пружины [m]
LF = 1  # Плечо приложения внешней силы [m]
k = 1  # Жёсткость пружины [N/m]
M = 2  # Масса груза [kg]
Mu1 = 0.1  # Коэффициент демпфирования [kg/s]
f1, f2 = [0, 5]  # Диапазон частот (начальное и конечное значения) [Hz]
F = 1  # Величина приложенной силы [N]

# Свободные колебания маятника
phi10 = np.pi / 2  # Начальный угол отклонения [rad]
omega10 = 0  # Начальная угловая скорость [rad/s]


# Определение функции для решения уравнения
def free_oscillations(y: np.ndarray, t: float) -> np.ndarray:
    """Функция, получающая на вход вектор Y = (y, y'), а возвращающая Y' = (y', y'').
    Необходима для интегрирования.

    Parameters
    ----------
    y : np.ndarray
        Вектор Y
    t : float
        Независимая переменная, в частном случае время

    Returns
    -------
    np.ndarray
        Вектор Y'
    """
    return [
        y[1],
        -g / L1 * np.sin(y[0])
        - Mu1 * y[1] / (M * L1**2)
        - k
        / (M * L1)
        * (L1 + L2)
        * (
            1
            - L2
            / (np.sqrt((L2) ** 2 + 2 * L1 * (L1 + L2) * (1 - np.cos(y[0]))))
            * np.sin(y[0])
        ),
    ]


# Решение системы ДУ для свободных колебаний
tk = 200  # Промежуток интегрирования
t1 = np.linspace(0, tk, 1000)
y1 = odeint(free_oscillations, [phi10, omega10], t1)

# График зависимости угла маятника от времени
fig4 = plt.figure()
plt.subplot(2, 1, 1)
plt.plot(t1, y1[:, 0])
plt.title("Свободные колебания маятника")
plt.xlabel("Время")
plt.ylabel("Угол маятника")

# АЧХ маятника с приложенной к нему силой
F1 = 1
phi10 = np.pi / 2
omega10 = 0


# Определение функции для решения уравнения с приложенной силой
def forced_oscillations(y: np.ndarray, t: float) -> np.ndarray:
    """Функция, получающая на вход вектор Y = (y, y'), а возвращающая Y' = (y', y'').
    Необходима для интегрирования.

    Parameters
    ----------
    y : np.ndarray
        Вектор Y
    t : float
        Независимая переменная, в частном случае время

    Returns
    -------
    np.ndarray
        Вектор Y'
    """
    return [
        y[1],
        -g / L1 * np.sin(y[0])
        - Mu1 * y[1] / (M * L1**2)
        - k
        / (M * L1)
        * (L1 + L2)
        * (
            1
            - L2
            / (np.sqrt((L2) ** 2 + 2 * L1 * (L1 + L2) * (1 - np.cos(y[0]))))
            * np.sin(y[0])
        )
        + L1 * F1 * np.sin(0.01 * np.pi * t**2) * np.cos(y[0]),
    ]


# Решение системы ДУ для колебаний с приложенной силой
t2 = np.linspace(0, tk, 1000)
y2 = odeint(forced_oscillations, [phi10, omega10], t2)

# График зависимости угла маятника от времени с приложенной силой
plt.subplot(2, 1, 2)
plt.plot(t2, y2[:, 0])
plt.title("Колебания с приложенной силой")
plt.xlabel("Время")
plt.ylabel("Угол маятника с приложенной силой")
fig4.savefig("dz/oscillation_analysis/images/image_oscillation.png")

# Собственные частоты колебаний первого и второго маятника
Omega1 = np.sqrt(
    g / L1
    + k / (M * L1) * (L1 + L2) * (1 - L2 / (np.sqrt((L2) ** 2 + 2 * L1 * (L1 + L2))))
)

# Частота гармонической силы
T = np.arange(0, 1000, 0.01)
p = 0.005 * T

# АЧХ 1-го и 2-го маятника
A1 = (F1 / M) / np.sqrt(
    np.abs(Omega1**2 - (2 * np.pi * p) ** 2) ** 2
    + 4 * (Mu1 / 2 * M) ** 2 * (2 * np.pi * p) ** 2
)

# График АЧХ маятника
fig5 = plt.figure("АЧХ")
plt.plot(p, A1)
plt.title("АЧХ маятника")
plt.xlabel("Частота, Гц")
plt.ylabel("Амплитуда, м")
plt.show()
fig5.savefig("dz/oscillation_analysis/images/image_response.png")
freq = round(Omega1 / (2 * np.pi), ndigits=1)


# Создание документа


docx_template = DocxTemplate(
    "dz/oscillation_analysis/docs/automated_report_template.docx"
)

image_beam = InlineImage(
    docx_template, "dz/oscillation_analysis/images/image_beam.png", Cm(12)
)
image_pendulum = InlineImage(
    docx_template, "dz/oscillation_analysis/images/image_pendulum.png", Cm(5)
)
image_bending = InlineImage(
    docx_template, "dz/oscillation_analysis/images/image_bending.png", Cm(13)
)
image_oscillation = InlineImage(
    docx_template, "dz/oscillation_analysis/images/image_oscillation.png", Cm(15)
)
image_response = InlineImage(
    docx_template, "dz/oscillation_analysis/images/image_response.png", Cm(15)
)
context = {
    "title1": "Анализ собственных колебаний балки",
    "title2": "Анализ механической системы груза и пружины",
    "year": datetime.datetime.now().strftime("%Y"),
    "image_beam": image_beam,
    "image_pendulum": image_pendulum,
    "image_bending": image_bending,
    "image_oscillation": image_oscillation,
    "image_response": image_response,
    "EI": EI,
    "l1": l1,
    "l2": l2,
    "P": P,
    "f1": f1,
    "f2": f2,
    "g": g,
    "L1": L1,
    "L2": L2,
    "LF": LF,
    "M": M,
    "k": k,
    "Mu1": Mu1,
    "phi10": phi10 * 180 / np.pi,
    "omega10": omega10,
    "Ra": Ra,
    "Rb": Rb,
    "freq": freq,
    "F": F,
    "y": y[-1],
    "Natural_frequency": Natural_frequency,
}

docx_template.render(context)
docx_template.save("dz/oscillation_analysis/docs/generated_report.docx")
